/*
 * Copyright 2002-2017 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.springframework.messaging.simp.config;

import org.springframework.lang.Nullable;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.util.Assert;

/**
 * A registration class for customizing the properties of {@link ThreadPoolTaskExecutor}.
 *
 * @author Rossen Stoyanchev
 * @author Juergen Hoeller
 * @since 4.0
 */
public class TaskExecutorRegistration {

    private final ThreadPoolTaskExecutor taskExecutor;

    @Nullable
    private Integer corePoolSize;

    @Nullable
    private Integer maxPoolSize;

    @Nullable
    private Integer keepAliveSeconds;

    @Nullable
    private Integer queueCapacity;


    /**
     * Create a new {@code TaskExecutorRegistration} for a default
     * {@link ThreadPoolTaskExecutor}.
     */
    public TaskExecutorRegistration() {
        this.taskExecutor = new ThreadPoolTaskExecutor();
        this.taskExecutor.setCorePoolSize(Runtime.getRuntime().availableProcessors() * 2);
        this.taskExecutor.setAllowCoreThreadTimeOut(true);
    }

    /**
     * Create a new {@code TaskExecutorRegistration} for a given
     * {@link ThreadPoolTaskExecutor}.
     *
     * @param taskExecutor the executor to use
     */
    public TaskExecutorRegistration(ThreadPoolTaskExecutor taskExecutor) {
        Assert.notNull(taskExecutor, "ThreadPoolTaskExecutor must not be null");
        this.taskExecutor = taskExecutor;
    }


    /**
     * Set the core pool size of the ThreadPoolExecutor.
     * <p><strong>NOTE:</strong> The core pool size is effectively the max pool size
     * when an unbounded {@link #queueCapacity(int) queueCapacity} is configured
     * (the default). This is essentially the "Unbounded queues" strategy as explained
     * in {@link java.util.concurrent.ThreadPoolExecutor ThreadPoolExecutor}. When
     * this strategy is used, the {@link #maxPoolSize(int) maxPoolSize} is ignored.
     * <p>By default this is set to twice the value of
     * {@link Runtime#availableProcessors()}. In an application where tasks do not
     * block frequently, the number should be closer to or equal to the number of
     * available CPUs/cores.
     */
    public TaskExecutorRegistration corePoolSize(int corePoolSize) {
        this.corePoolSize = corePoolSize;
        return this;
    }

    /**
     * Set the max pool size of the ThreadPoolExecutor.
     * <p><strong>NOTE:</strong> When an unbounded
     * {@link #queueCapacity(int) queueCapacity} is configured (the default), the
     * max pool size is effectively ignored. See the "Unbounded queues" strategy
     * in {@link java.util.concurrent.ThreadPoolExecutor ThreadPoolExecutor} for
     * more details.
     * <p>By default this is set to {@code Integer.MAX_VALUE}.
     */
    public TaskExecutorRegistration maxPoolSize(int maxPoolSize) {
        this.maxPoolSize = maxPoolSize;
        return this;
    }

    /**
     * Set the time limit for which threads may remain idle before being terminated.
     * If there are more than the core number of threads currently in the pool,
     * after waiting this amount of time without processing a task, excess threads
     * will be terminated. This overrides any value set in the constructor.
     * <p>By default this is set to 60.
     */
    public TaskExecutorRegistration keepAliveSeconds(int keepAliveSeconds) {
        this.keepAliveSeconds = keepAliveSeconds;
        return this;
    }

    /**
     * Set the queue capacity for the ThreadPoolExecutor.
     * <p><strong>NOTE:</strong> when an unbounded {@code queueCapacity} is configured
     * (the default), the core pool size is effectively the max pool size. This is
     * essentially the "Unbounded queues" strategy as explained in
     * {@link java.util.concurrent.ThreadPoolExecutor ThreadPoolExecutor}. When
     * this strategy is used, the {@link #maxPoolSize(int) maxPoolSize} is ignored.
     * <p>By default this is set to {@code Integer.MAX_VALUE}.
     */
    public TaskExecutorRegistration queueCapacity(int queueCapacity) {
        this.queueCapacity = queueCapacity;
        return this;
    }


    protected ThreadPoolTaskExecutor getTaskExecutor() {
        if (this.corePoolSize != null) {
            this.taskExecutor.setCorePoolSize(this.corePoolSize);
        }
        if (this.maxPoolSize != null) {
            this.taskExecutor.setMaxPoolSize(this.maxPoolSize);
        }
        if (this.keepAliveSeconds != null) {
            this.taskExecutor.setKeepAliveSeconds(this.keepAliveSeconds);
        }
        if (this.queueCapacity != null) {
            this.taskExecutor.setQueueCapacity(this.queueCapacity);
        }
        return this.taskExecutor;
    }

}
